package com.xcj.springai.controller;

import com.alibaba.cloud.ai.graph.*;
import com.alibaba.cloud.ai.graph.async.AsyncGenerator;
import com.alibaba.cloud.ai.graph.checkpoint.config.SaverConfig;
import com.alibaba.cloud.ai.graph.checkpoint.constant.SaverConstant;
import com.alibaba.cloud.ai.graph.checkpoint.savers.MemorySaver;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.StateSnapshot;
import com.xcj.springai.config.GraphProcess;
import com.xcj.springai.dto.ResumeDto;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicReference;


@RestController
@RequestMapping("/graph/human")
@Slf4j
public class GraphHumanController {

    private static final Logger logger = LoggerFactory.getLogger(GraphHumanController.class);

    private final CompiledGraph compiledGraph;

    @Autowired
    public GraphHumanController(@Qualifier("humanGraph") StateGraph stateGraph) throws GraphStateException {
        SaverConfig saverConfig = SaverConfig.builder().register(SaverConstant.MEMORY, new MemorySaver()).build();
        this.compiledGraph = stateGraph
                .compile(CompileConfig.builder().saverConfig(saverConfig).interruptBefore("humanfeedback").build());    }

    @GetMapping(value = "/expand", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> expand(@RequestParam(value = "query", defaultValue = "你好，很高兴认识你，能简单介绍一下自己吗？", required = false) String query,
                                                @RequestParam(value = "expander_number", defaultValue = "3", required = false) Integer expanderNumber,
                                                @RequestParam(value = "threadId", defaultValue = "yingzi", required = false) String threadId
    ) throws GraphRunnerException {
        RunnableConfig runnableConfig = RunnableConfig.builder().threadId(threadId).build();
        Map<String, Object> objectMap = new HashMap<>();
        objectMap.put("query", query);
        objectMap.put("expander_number", expanderNumber);

        GraphProcess graphProcess = new GraphProcess(this.compiledGraph);
        Sinks.Many<ServerSentEvent<String>> sink = Sinks.many().unicast().onBackpressureBuffer();
        AsyncGenerator<NodeOutput> resultFuture = compiledGraph.stream(objectMap, runnableConfig);
        graphProcess.processStream(resultFuture, sink);

        GraphRepresentation graphRepresentation = compiledGraph.getGraph(GraphRepresentation.Type.PLANTUML);
        System.out.println("\n\n");
        System.out.println(graphRepresentation.content());
        System.out.println("\n\n");
        return sink.asFlux()
                .doOnCancel(() -> logger.info("Client disconnected from stream"))
                .doOnError(e -> logger.error("Error occurred during streaming", e));
    }

    @PostMapping(value = "/resume", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> resume(@RequestBody ResumeDto resumeDto) throws GraphRunnerException {
        // 构造 RunnableConfig，使用客户端传入的线程ID，用于区分图执行的上下文线程
        RunnableConfig runnableConfig = RunnableConfig.builder()
                .threadId(resumeDto.getThreadId())
                .build();

        // 从 compiledGraph 获取对应线程的历史状态快照（保存了之前图执行的状态）
        StateSnapshot stateSnapshot = this.compiledGraph.getState(runnableConfig);
        OverAllState state = stateSnapshot.state(); // 取得整体状态对象

        // 从请求体中获取人工反馈数据
        Map<String, Object> node = resumeDto.getHumanFeedback();
        String feedback = resumeDto.getHumanFeedback().toString();
        state.withResume();

        OverAllState.HumanFeedback humanFeedback = new OverAllState.HumanFeedback(node, feedback);
        state.withHumanFeedback(humanFeedback);
        Map<String, Object> update = new HashMap<>();
        update.put("feedback", resumeDto.getHumanFeedback());
        state.updateState(update);
        state.withHumanFeedback(new OverAllState.HumanFeedback(update, ""));

        Sinks.Many<ServerSentEvent<String>> sink = Sinks.many().unicast().onBackpressureBuffer();
        GraphProcess graphProcess = new GraphProcess(this.compiledGraph);
        AsyncGenerator<NodeOutput> stream = compiledGraph.streamFromInitialNode(state, runnableConfig);
        graphProcess.processStream(stream, sink);
        logger.info("控制器层 Resuming graph with threadId: {}", resumeDto.getThreadId());
        logger.info("控制器层 State object hash: {}", System.identityHashCode(state));
        logger.info("控制器层 Injected human feedback: {}", resumeDto.getHumanFeedback());
        logger.info("控制器层 Before stream: humanFeedback = {}", state.humanFeedback());

        // 返回响应式事件流，支持客户端断开和异常的日志处理
        return sink.asFlux()
                .doOnCancel(() -> logger.info("Client disconnected from stream"))
                .doOnError(e -> logger.error("Error occurred during streaming", e));
    }



}
